#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版 SSH 蜜罐处理器（基于 asyncssh）
与现有 ProtocolHandler 集成：继承自 ProtocolHandler，并调用 self.persist(entry) 保存事件。
功能：
 - 完整 SSH 握手（KEX、host key）
 - 认证点位（password / publickey）：记录凭证尝试
 - 交互 shell（PTY）与 exec：记录命令与响应
 - 会话元数据（session_id, client_version, remote ip/port, start/end/duration）
 - 可配置：是否接受密码、公钥；欢迎词、prompt、伪文件内容等
注意：不要在生产网络直接暴露未隔离的蜜罐实例，务必放在隔离网络。
"""
from __future__ import annotations

import asyncio
import asyncssh
import base64
import uuid
import logging
from datetime import datetime, timezone
from typing import Dict, Any, Optional

from .base_handler import ProtocolHandler
from .config import get_ssh_config
from .logger import op_logger
from multihoneypot.attacklogger import attack_logger, AttackType, ProtocolType

def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()

class HoneySSHServer(asyncssh.SSHServer):
    """
    asyncssh 的 server-side 回调：处理连接级事件与认证回调委托给 handler。
    """

    def __init__(self, handler: "SSHHandler", peername):
        self.handler = handler
        self.peername = peername  # (ip, port)

    def connection_made(self, conn):
        # conn: asyncssh.SSHServerConnection
        client_version = getattr(conn, "client_version", None)
        op_logger.info("connection_made peer=%s client_version=%s", self.peername, client_version)

    def connection_lost(self, exc):
        op_logger.info("connection_lost peer=%s exc=%s", self.peername, exc)

    def begin_auth(self, username: str) -> bool:
        # True 表示需要认证（将触发 validate_password / validate_public_key）
        return True

    async def validate_password(self, username: str, password: str) -> bool:
        # 记录认证尝试到攻击日志
        attack_logger.log_auth_attempt(
            source_ip=self.peername[0],
            destination_port=self.handler.port,
            protocol=ProtocolType.SSH,
            username=username,
            password=password
        )

        # 记录尝试（避免在日志中直接打印明文密码，请注意权限）
        await self.handler._record_auth_attempt({
            "method": "password",
            "username": username,
            "password": password,  # 如果合规/敏感，后续可替换为 hash
            "remote_ip": self.peername[0],
            "remote_port": self.peername[1],
            "time": now_iso(),
        })
        # 根据配置决定是否允许（蜜罐策略）
        # return bool(self.handler.config.get("accept_password", False))
        allowed: Dict[str, str] = self.handler.config.get("allowed_credentials", {})
        if username in allowed and allowed[username] == password:
            return True
        return False

    async def validate_public_key(self, username: str, key: asyncssh.PKey) -> bool:
        """
        验证SSH公钥认证请求

        Args:
            username (str): 尝试认证的用户名
            key (asyncssh.PKey): 客户端提供的公钥对象

        Returns:
            bool: 如果接受公钥认证则返回True，否则返回False
        """
        fp = key.get_fingerprint()

        # 记录认证尝试到攻击日志
        attack_logger.log_auth_attempt(
            source_ip=self.peername[0],
            destination_port=self.handler.port,
            protocol=ProtocolType.SSH,
            username=username,
            pubkey_fingerprint=fp
        )

        await self.handler._record_auth_attempt({
            "method": "publickey",
            "username": username,
            "pubkey_fingerprint": fp,
            "pubkey_type": getattr(key, "get_type", lambda: "unknown")(),
            "remote_ip": self.peername[0],
            "remote_port": self.peername[1],
            "time": now_iso(),
        })
        return bool(self.handler.config.get("accept_publickey", False))

class HoneySSHSession(asyncssh.SSHServerSession):
    """
    处理单个 session（shell/exec/pty）
    """

    def __init__(self, handler: "SSHHandler", session_meta: Dict[str, Any]):
        self.handler = handler
        self.session_meta = session_meta
        self._chan = None
        self._buffer = ""
        self.pty_info = None
        self._closed = False

    def connection_made(self, chan):
        self._chan = chan
        op_logger.debug("session connection_made session_id=%s", self.session_meta.get("session_id"))

    def connection_lost(self, exc):
        op_logger.debug("session connection_lost session_id=%s exc=%s", self.session_meta.get("session_id"), exc)

    def session_started(self):
        # callback when session established
        pass

    def pty_requested(self, term, width, height, pxwidth, pxheight, modes):
        self.pty_info = {"term": term, "width": width, "height": height}
        return True

    def shell_requested(self):
        # 启动交互 shell
        asyncio.create_task(self._run_shell())
        return True

    def exec_requested(self, command):
        # 记录命令执行到攻击日志
        if self.session_meta.get("remote_ip"):
            attack_logger.log_command_execution(
                source_ip=self.session_meta["remote_ip"],
                destination_port=self.session_meta.get("destination_port", 22),
                protocol=ProtocolType.SSH,
                command=command,
                command_type="exec"
            )

        # 记录 exec 命令并立即返回伪结果后退出
        asyncio.create_task(self.handler._record_command(self.session_meta, command, is_exec=True))
        # 使用 handler 的 fake_exec_response 生成结果
        out = self.handler.fake_exec_response(command)
        try:
            # write 可能在不同线程上下文，使用 chan.write
            if self._chan:
                self._chan.write(out + "\n")
                self._chan.exit(0)
        except Exception:
            op_logger.exception("exec response error")
        return True

    async def _run_shell(self):
        """
        读取 channel 数据（byte），按行解析命令并响应。
        同时记录原始字节（base64）与解析后的命令事件。
        """
        prompt = self.handler.shell_prompt()
        welcome = self.handler.welcome_message()
        try:
            if self._chan is None:
                return
            # 写欢迎词和 prompt
            if welcome:
                self._chan.write(welcome + "\r\n")
            if prompt:
                self._chan.write(prompt)
            # 读取循环
            while True:
                data = await self._chan.read(1024)
                if data is None:
                    break
                # data 是 bytes
                try:
                    decoded = data.decode(errors="ignore")
                except Exception:
                    decoded = str(data)
                # 记录原始交互片段
                await self.handler._record_interaction(self.session_meta, data)
                # 拼接缓冲并按换行拆分以得到完整命令
                self._buffer += decoded
                if "\n" in self._buffer or "\r" in self._buffer:
                    # 取第一条命令（保守）
                    line, _, rest = self._buffer.partition("\n")
                    line = line.rstrip("\r")
                    self._buffer = rest
                    cmd = line.strip()

                    # 记录命令执行到攻击日志
                    if self.session_meta.get("remote_ip") and cmd:
                        attack_logger.log_command_execution(
                            source_ip=self.session_meta["remote_ip"],
                            destination_port=self.session_meta.get("destination_port", 22),
                            protocol=ProtocolType.SSH,
                            command=cmd,
                            command_type="shell"
                        )

                    # 记录命令事件
                    await self.handler._record_command(self.session_meta, cmd, is_exec=False)
                    # 特殊命令处理
                    if cmd in ("exit", "logout", "quit"):
                        self._chan.write("logout\r\n")
                        self._chan.exit(0)
                        break
                    # 生成伪输出并写回
                    resp = self.handler.fake_exec_response(cmd)
                    if resp:
                        self._chan.write(resp + "\r\n")
                    # 再写 prompt
                    if prompt:
                        self._chan.write(prompt)
                else:
                    # 未到换行，继续读取
                    pass
        except Exception as e:
            op_logger.exception("shell loop exception session=%s: %s", self.session_meta.get("session_id"), e)
        finally:
            self._closed = True

class SSHHandler(ProtocolHandler):
    """
    可直接替换你现有的 SSH Handler：
      - 保持与 ProtocolHandler.persist(entry) 的兼容（调用 await self.persist(entry)）
      - 配置源自 get_ssh_config(config)
    """

    def __init__(self, host: str, port: int, log_file: str, config=None):
        super().__init__("ssh", host, port, log_file)
        self.config: Dict[str, Any] = {}
        if config:
            self.config.update(get_ssh_config(config))
        # 默认值
        self.config.setdefault("banner", "SSH-2.0-OpenSSH_7.9")
        self.config.setdefault("accept_password", False)
        self.config.setdefault("accept_publickey", False)
        self.config.setdefault("shell_prompt", "[root@server ~]$ ")
        self.config.setdefault("welcome_message", "Welcome to Ubuntu 18.04.6 LTS (GNU/Linux 4.15.0-xyz)")
        self.config.setdefault("host_key_path", None)  # 若配置了 path 则读取私钥文件
        self.config.setdefault("allowed_credentials", {"root": "123456", "admin": "adminpwd"})

        self._server = None
        self._active_sessions = {}  # session_id -> meta

    def shell_prompt(self) -> str:
        return str(self.config.get("shell_prompt", ""))

    def welcome_message(self) -> str:
        return str(self.config.get("welcome_message", ""))

    async def _record_auth_attempt(self, info: Dict[str, Any]):
        # 标准化保存
        entry = {
            "event": "auth_attempt",
            "time": now_iso(),
            **info
        }
        try:
            await self.persist(entry)
        except Exception:
            op_logger.exception("persist auth attempt failed")

    async def _record_command(self, session_meta: Dict[str, Any], command: str, is_exec: bool = False):
        entry = {
            "event": "command",
            "time": now_iso(),
            "session_id": session_meta.get("session_id"),
            "username": session_meta.get("username"),
            "command": command,
            "is_exec": bool(is_exec),
            "remote_ip": session_meta.get("remote_ip"),
            "remote_port": session_meta.get("remote_port"),
        }
        try:
            await self.persist(entry)
        except Exception:
            op_logger.exception("persist command failed")

    async def _record_interaction(self, session_meta: Dict[str, Any], raw_bytes: bytes):
        entry = {
            "event": "interaction",
            "time": now_iso(),
            "session_id": session_meta.get("session_id"),
            "username": session_meta.get("username"),
            "raw_base64": base64.b64encode(raw_bytes).decode("ascii"),
            "remote_ip": session_meta.get("remote_ip"),
            "remote_port": session_meta.get("remote_port"),
        }
        try:
            await self.persist(entry)
        except Exception:
            op_logger.exception("persist interaction failed")

    def fake_exec_response(self, cmd: str) -> str:
        # 伪造命令输出（可扩展成模板或脚本）
        cmd = (cmd or "").strip()
        if cmd in ("ls", "ls -la"):
            return "total 8\ndrwxr-xr-x 2 root root 4096 Sep 20 12:00 .\n-rw-r--r-- 1 root root   23 Sep 20 12:00 README"
        if cmd.startswith("cat "):
            return "This is a fake file content"
        if cmd in ("whoami",):
            return "root"
        if cmd in ("pwd",):
            return "/root"
        if cmd.startswith("uname") or cmd.startswith("uname -a"):
            return "Linux server 4.15.0-xyz #1 SMP Thu Sep 10 12:00:00 UTC 2020 x86_64 GNU/Linux"
        if cmd.startswith("wget ") or cmd.startswith("curl "):
            return "Downloading... saved to /tmp/fakefile"
        if cmd == "":
            return ""
        return f"bash: {cmd}: command not found"

    async def start(self):
        """
        启动 asyncssh 服务。会生成临时 host key（除非配置 host_key_path）。
        """
        # 读取配置 host key 或生成临时 host key
        host_key_path = self.config.get("host_key_path")
        server_host_keys = None
        if host_key_path:
            try:
                server_host_keys = [host_key_path]
            except Exception:
                op_logger.exception("failed to use host_key_path %s", host_key_path)
                server_host_keys = None

        if not server_host_keys:
            # 生成临时 host key（RSA）
            try:
                # asyncssh.generate_private_key 返回 PKey，create_server 接受 PKey 对象
                temp_key = asyncssh.generate_private_key("ssh-rsa")
                server_host_keys = [temp_key]
                op_logger.info("generated ephemeral host key")
            except Exception:
                op_logger.exception("failed to generate ephemeral host key")
                server_host_keys = None  # let asyncssh decide (rare)

        # 启动 asyncssh server
        op_logger.info("starting asyncssh server on %s:%s", self.host, self.port)
        try:
            server = await asyncssh.create_server(
                lambda: HoneySSHServer(self, None),
                host=self.host,
                port=self.port,
                server_host_keys=server_host_keys,
                # 下面参数可以根据需要调整
                allow_scp=False,
                # 如果要支持 SFTP，你需要传入 sftp_factory
            )
            self._server = server
            op_logger.info("asyncssh server started listen=%s:%s", self.host, self.port)
            # asyncssh.create_server 返回一个 Server 对象；保持运行由外部 event loop 管理
        except Exception as e:
            op_logger.exception("failed to start asyncssh server: %s", e)
            raise

    async def stop(self):
        if self._server:
            try:
                self._server.close()
                await self._server.wait_closed()
                op_logger.info("ssh server closed")
            except Exception:
                op_logger.exception("error while stopping ssh server")
